import os
from PIL import Image


def make_yolo_format(src_path:str,dst_path:str):
    with open(f"{src_path}.cat", 'r') as file:
        numbers = [int(num) for line in file for num in line.split()]
    if numbers[0]!=9:
        return False
    numbers=numbers[1:]

    with Image.open(src_path) as img:
        width, height = img.size
    
    numbers=[numbers[i]/(width if i%2==0 else height) for i in range(len(numbers)) ]
    x_list=[numbers[i] for i in range(len(numbers)) if i % 2 == 0]
    y_list=[numbers[i] for i in range(len(numbers)) if i % 2 != 0]
    x_min=min(x_list)
    x_max=max(x_list)
    y_min=min(y_list)
    y_max=max(y_list)

    res=[0,(sum(x_list)+2*x_list[2])/(len(x_list)+2),(sum(y_list)+2*x_list[2])/(len(y_list)+2),1.1*(x_max-x_min),1.1*(y_max-y_min),*numbers]
    if min(res)<0 or max(res)>1:
        return False
    with open(dst_path, 'w') as file:
    # 将整数列表转换为字符串，并用空格分割
        file.write(' '.join(map(str, res)))
    return True

def prepare(src_dir:str,dst_dir:str):
    prefix=os.path.basename(src_dir)
    src_dir=os.path.abspath(src_dir)
    for filename in os.listdir(src_dir):
        if filename.endswith(".jpg"):
            basename=os.path.splitext(filename)[-2]
            src_path=f"{src_dir}/{filename}"
            img_dst_path=f"{dst_dir}/images/{prefix}_{basename}.jpg"
            txt_dst_path=f"{dst_dir}/labels/{prefix}_{basename}.txt"
            if make_yolo_format(src_path,txt_dst_path) and not os.path.lexists(img_dst_path):
                os.symlink(src_path,img_dst_path)

def clean(dir:str):
    for filename in os.listdir(dir):
        try:
            os.remove(f"{dir}/{filename}")
        except:
            pass

if __name__=="__main__":
    train_src_dirs=["datasets/cat_pose/CAT_00","datasets/cat_pose/CAT_01","datasets/cat_pose/CAT_02"]
    train_dst_dir="datasets/assets/train"

    val_src_dirs=["datasets/cat_pose/CAT_05","datasets/cat_pose/CAT_06"]
    val_dst_dir="datasets/assets/val"

    clean(f"{train_dst_dir}/images")
    clean(f"{train_dst_dir}/labels")
    clean(f"{val_dst_dir}/images")
    clean(f"{val_dst_dir}/labels")

    for dir in train_src_dirs:
        prepare(dir,train_dst_dir)

    for dir in val_src_dirs:
        prepare(dir,val_dst_dir)
